from typing import Dict, List

from langchain.prompts import FewShotPromptTemplate, PromptTemplate
from langchain.chains import create_retrieval_chain
from langchain_core.messages import SystemMessage, HumanMessage, ToolMessage, AIMessage
from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate, HumanMessagePromptTemplate, \
    AIMessagePromptTemplate, SystemMessagePromptTemplate
from langchain_core.tools import tool
from langchain_experimental.tabular_synthetic_data.prompts import (
    SYNTHETIC_FEW_SHOT_PREFIX,
    SYNTHETIC_FEW_SHOT_SUFFIX,
)
from langchain_openai import ChatOpenAI
from pydantic import BaseModel
from langchain_community.embeddings import FastEmbedEmbeddings
from langchain_chroma import Chroma

vectorstore = Chroma(
    collection_name="full_documents", embedding_function=FastEmbedEmbeddings(),
    persist_directory="../data/"
)


@tool(response_format="content_and_artifact")
def retrieve(query: str):
    """Retrieve information related to a query."""
    retrieved_docs = vectorstore.similarity_search(query, k=2)
    serialized = "\n\n".join(
        (f"Source: {doc.metadata}\n" f"Content: {doc.page_content}")
        for doc in retrieved_docs
    )

    return serialized, retrieved_docs


tools = [retrieve]

question_system_prompt = f'''
    你是一个数据增强的专家
    question的格式如下：
    格式相同，但性别，年龄，日期，入睡时间，起床时间不一样，D,S,R,A不一样。
'''

answer_system_prompt = f'''
    你是一名睡眠科学专家，擅长从可穿戴设备数据中解读睡眠质量，能结合睡眠医学理论与用户行为模式提供针对性建议。
根据昨日数据给出睡眠质量评分（10分制）
对比其他时间段，分析昨日的深睡、浅睡、REM、清醒、入睡时间、起床时间，格式多样
给出合理性建议，随机两个项给出建议，风格多样：
要求：
    - 用中文回复，400-500字符
    - 如果是女性，保持语气温柔、礼貌，让对方感受到你的真诚和关心、如果是男性，真诚直接
    - 不要出现*这个符号
    - 建议的内容要和用户的睡眠数据相关，不要胡编乱造
格式：
    - 评分：...
    - 对比分析：...
    - 建议：...
    '''

question_examples = [
    {
        "question": '''
    女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
    ''',
    },
    {
        "question": '''
    女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
    ''',
    },
    {
        "question": '''
    女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
    '''
    }
]

answer_examples = [
    {
        "input": '''
女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
        ''',
        "output": '''
        您好！根据您提供的睡眠数据，我为您进行以下分析：

评分：6.5分。昨日睡眠总时长尚可，但深睡比例偏低，REM睡眠时间稍长，清醒时间较短，整体睡眠结构有待优化。

对比分析：

*   深睡：昨日深睡时间195分钟，远高于其他几日的平均值（约118分钟），说明您在深睡阶段得到了较好的休息。
*   浅睡：昨日浅睡时间200分钟，与平均值接近，波动不大。
*   REM：昨日REM睡眠95分钟，高于平均水平（约80分钟），REM时间过长可能导致白天疲惫。
*   清醒：昨日清醒时间10分钟，是这几天中最少的，说明您在夜间保持了较好的睡眠连贯性。
*   入睡时间：昨日入睡时间为22:50，与以往数据接近。
*   起床时间：昨日起床时间为07:40，比其他日子晚起，这可能是由于深睡时间更长。

建议：

1.  鉴于您昨日深睡时间较长，可以尝试规律作息，保持每天的入睡和起床时间一致，这有助于稳定您的睡眠周期，提高整体睡眠质量。另外，深睡时间过长可能导致白天犯困，建议您关注白天的精神状态，看看是否有嗜睡的情况。
2.  由于REM睡眠时间偏长，建议您在睡前一小时避免剧烈运动和过量饮水，晚餐尽量清淡，避免影响您的REM睡眠。

希望我的分析对您有所帮助，祝您拥有一个美好的睡眠！
        '''
    },
    {
        "input": '''
女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
                ''',
        "output": '''
                您好，基于您提供的睡眠数据，我将为您进行详细的分析：

- 评分：7.5分。您的睡眠总时长相对充足，但深睡时间有待提升，清醒时间偏短，整体睡眠质量尚可。

- 对比分析：
    *   深睡：昨日的深睡时间为195分钟，明显高于其他几天，说明昨日的深睡质量不错。
    *   浅睡：昨日浅睡200分钟，与前几天相比，浅睡时间大致持平，没有太大波动。
    *   REM：昨日REM睡眠95分钟，与11月12日相同，REM睡眠时间充足，对巩固记忆和情绪调节有益。
    *   清醒：昨日清醒时间仅10分钟，是所有数据中最少的，这表明您在睡眠过程中保持了较好的连续性。
    *   入睡时间：昨日22:50入睡，与11月12日相同，其他天都是23点左右入睡。
    *   起床时间：昨日07:40起床，是所有记录中最晚的一次，这可能与您前一晚的睡眠质量有关。

- 建议：
    1.  为了更好地维持您的睡眠质量，建议您尽量保持规律的作息时间，包括入睡和起床时间。即使在周末，也尽量不要打破这种规律，这有助于调节您的生物钟。
    2.  深睡时间对于身体的修复和恢复至关重要。您可以尝试在睡前进行放松活动，如阅读、听轻音乐或进行舒缓的瑜伽。避免睡前饮用咖啡因和酒精饮料，这些物质可能会影响您的深睡时间。
                '''
    },
    {
        "input": '''
女性 50岁
11/10 2320-0710 D110 S215 R80 A25
11/11 0015-0800 D095 S230 R70 A35
11/12 2250-0640 D125 S200 R95 A20
11/13 2330-0725 D105 S240 R75 A30
11/14 2245-0630 D130 S195 R100 A15
11/15 2310-0715 D115 S220 R85 A25
昨日   2250-0740 D195 S200 R95 A10
D=深睡(min) / S=浅睡(min)
R=REM(min) / A=清醒(min)
时间采用24小时制简写（如2320=23:20）
                ''',
        "output": '''
                好的，根据您提供的睡眠数据，我将为您进行详细的分析：

评分：7.5分

对比分析：
- 入睡时间：昨日22:50，与前几日相比，入睡时间相对较早。
- 起床时间：昨日07:40，与前几日相比，起床时间略晚。
- 深睡：昨日深睡时间195分钟，是这段时间最长的。这表明您的身体在昨日得到了更充分的休息。
- 浅睡：昨日浅睡时间200分钟，与前几日相比，基本持平。
- REM：昨日REM睡眠时间95分钟，为这段时间最高，REM睡眠对记忆巩固和情绪调节至关重要。
- 清醒：昨日清醒时间10分钟，是这段时间最低，说明夜间睡眠较为稳定，很少出现醒来。

建议：
1.  **保持规律的作息**：您昨日的睡眠时长相对较长，深度睡眠时间也很理想。我建议您尽量保持规律的作息习惯，在您感觉舒适的时间入睡和起床。固定睡眠时间有助于调节您的生物钟，让您更容易入睡并获得高质量的睡眠。
2.  **营造舒适的睡眠环境**：您昨日的深睡和REM睡眠数据都很好。为了保持这种高质量的睡眠，您可以尝试改善您的睡眠环境。确保卧室黑暗、安静、凉爽，避免睡前使用电子设备，都有助于提升睡眠质量。
                '''
    },
]

question_prompt = PromptTemplate.from_template(template="{question}")
answer_prompt = PromptTemplate.from_template(template="{input},{output}")

question_prompt_template = FewShotPromptTemplate(
    prefix=SYNTHETIC_FEW_SHOT_PREFIX,
    examples=question_examples,
    suffix=SYNTHETIC_FEW_SHOT_SUFFIX,
    input_variables=["subject", "extra"],
    example_prompt=question_prompt,
)

answer_few_shot_prompt = FewShotChatMessagePromptTemplate(
    input_variables=["input"],
    examples=answer_examples,
    example_prompt=(
        HumanMessagePromptTemplate.from_template("{input}")
        + AIMessagePromptTemplate.from_template("{output}")
    ),
)

class QuestionSchema(BaseModel):
    question: str


model = ChatOpenAI(
    temperature=1,
    model="gpt-4o-mini",
    cache=False
)

question_chain = (
        question_prompt_template | model.with_structured_output(schema=QuestionSchema)
)

model_with_tools = model.bind_tools(tools)

questions = []

for i in range(3):
    question_schema = question_chain.invoke({"subject": "question_answer", "extra": question_system_prompt})
    print(question_schema)
    questions.append(question_schema)

# final_answer_prompt = ChatPromptTemplate.from_messages(
#         [
#             SystemMessage(content=answer_system_prompt),
#             answer_few_shot_prompt,
#             HumanMessage(content="{input}"),
#         ]
#     )

final_answer_prompt = (
    SystemMessagePromptTemplate.from_template(
        answer_system_prompt
    )
    + HumanMessagePromptTemplate.from_template("{input}")
)



for question in questions:
    print("一个完整流程")

    answer_chain = (final_answer_prompt | model_with_tools)
    result = answer_chain.invoke({"input":question.question})
    print(result)
    messages = [
        result
    ]

    for tool_call in result.tool_calls:
        selected_tool = {"retrieve": retrieve}[tool_call["name"].lower()]
        tool_result = selected_tool.invoke(tool_call)
        print(tool_result)
        messages.append(ToolMessage(
            content=tool_result,
            tool_call_id=tool_call["id"],
            name=tool_call["name"],
            args=tool_call["args"]
        ))

    ai_tool_msg = model_with_tools.invoke(messages)
    print(ai_tool_msg)